"""
仓位管理模块

管理回测过程中的仓位变化和资金使用。

作者: AI Assistant
版本: 1.0.0
日期: 2025-01-06
"""

from typing import Dict, List
import pandas as pd


class PositionManager:
    """
    仓位管理器
    
    负责跟踪和管理回测过程中的持仓情况。
    
    属性:
        initial_capital: 初始资金
        current_cash: 当前现金
        current_position: 当前持仓
        position_history: 持仓历史记录
    """
    
    def __init__(self, initial_capital: float):
        """
        初始化仓位管理器
        
        参数:
            initial_capital: 初始资金
        """
        self.initial_capital = initial_capital
        self.current_cash = initial_capital
        self.current_position = {
            'quantity': 0,  # 持仓数量
            'cost': 0.0,    # 持仓成本
            'value': 0.0,   # 持仓市值
        }
        self.position_history = []
    
    def update_position(
        self,
        quantity: int,
        price: float,
        cash: float,
        date
    ):
        """
        更新持仓信息
        
        参数:
            quantity: 持仓数量
            price: 当前价格
            cash: 当前现金
            date: 日期
        """
        self.current_position['quantity'] = quantity
        self.current_position['value'] = quantity * price
        self.current_cash = cash
        
        # 记录历史
        self.position_history.append({
            'date': date,
            'quantity': quantity,
            'price': price,
            'value': quantity * price,
            'cash': cash,
            'total': cash + quantity * price,
        })
    
    def get_total_value(self, current_price: float) -> float:
        """
        获取总资产价值
        
        参数:
            current_price: 当前股票价格
            
        返回:
            float: 总资产（现金+持仓市值）
        """
        position_value = self.current_position['quantity'] * current_price
        return self.current_cash + position_value
    
    def get_position_info(self) -> Dict:
        """
        获取当前持仓信息
        
        返回:
            Dict: 持仓信息字典
        """
        return {
            '持仓数量': self.current_position['quantity'],
            '持仓成本': self.current_position['cost'],
            '持仓市值': self.current_position['value'],
            '当前现金': self.current_cash,
            '总资产': self.current_cash + self.current_position['value'],
        }
    
    def get_position_history_df(self) -> pd.DataFrame:
        """
        获取持仓历史DataFrame
        
        返回:
            pd.DataFrame: 持仓历史数据
        """
        if not self.position_history:
            return pd.DataFrame()
        return pd.DataFrame(self.position_history).set_index('date')

